from ultralytics import YOLO
import cv2

# Load a model
#model = YOLO("yolo11n.pt")

def model_train():
    # Train the model
    train_results = model.train(
        data="coco8.yaml",  # path to dataset YAML
        epochs=100,  # number of training epochs
        imgsz=640,  # training image size
        device="cpu",  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
    )

def model_val():
    # Evaluate model performance on the validation set
    metrics = model.val()


def model_inference(image):
    # Perform object detection on an image
    results = model(image)
    print(results)

def model_export():
    # Export the model to ONNX format
    return  model.export(format="onnx")  # return path to exported model

def predict(chosen_model, img, classes=[], conf=0.5):
    if classes:
        results = chosen_model.predict(img, classes=classes, conf=conf)
    else:
        results = chosen_model.predict(img, conf=conf)
    return results

def predict_and_detect(chosen_model, img, classes=[], conf=0.5, rectangle_thickness=2, text_thickness=1):
    results = predict(chosen_model, img, classes, conf=conf)
    for result in results:
        for box in result.boxes:
            cv2.rectangle(img, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),
                          (int(box.xyxy[0][2]), int(box.xyxy[0][3])), (255, 0, 0), rectangle_thickness)
            cv2.putText(img, f"{result.names[int(box.cls[0])]}",
                        (int(box.xyxy[0][0]), int(box.xyxy[0][1]) - 10),
                        cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), text_thickness)
    return img, results

model = YOLO("yolo11m.pt")
# read the image
filename = "IMG_3969"
filein = "/home/v/kn/dev/ai/testdata/" + filename + ".jpg"
fileout = "/home/v/kn/dev/ai/testdata/" + filename + "-out.jpg"
image = cv2.imread(filein)
result_img, _ = predict_and_detect(model, image, classes=[], conf=0.5)
cv2.imwrite(fileout, result_img)

#model_inference("/home/v/kn/dev/ai/testdata/cabinet.jpg")
